{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a1199016",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: optimum==1.16.1 in /opt/conda/lib/python3.10/site-packages (1.16.1)\n",
      "Requirement already satisfied: coloredlogs in /opt/conda/lib/python3.10/site-packages (from optimum==1.16.1) (15.0.1)\n",
      "Requirement already satisfied: sympy in /opt/conda/lib/python3.10/site-packages (from optimum==1.16.1) (1.12)\n",
      "Requirement already satisfied: transformers>=4.26.0 in /opt/conda/lib/python3.10/site-packages (from transformers[sentencepiece]>=4.26.0->optimum==1.16.1) (4.36.0)\n",
      "Requirement already satisfied: torch>=1.9 in /opt/conda/lib/python3.10/site-packages (from optimum==1.16.1) (2.2.0)\n",
      "Requirement already satisfied: packaging in /opt/conda/lib/python3.10/site-packages (from optimum==1.16.1) (21.3)\n",
      "Requirement already satisfied: numpy in /opt/conda/lib/python3.10/site-packages (from optimum==1.16.1) (1.23.4)\n",
      "Requirement already satisfied: huggingface-hub>=0.8.0 in /opt/conda/lib/python3.10/site-packages (from optimum==1.16.1) (0.20.1)\n",
      "Requirement already satisfied: datasets in /opt/conda/lib/python3.10/site-packages (from optimum==1.16.1) (2.3.2)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.10/site-packages (from huggingface-hub>=0.8.0->optimum==1.16.1) (3.13.1)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /opt/conda/lib/python3.10/site-packages (from huggingface-hub>=0.8.0->optimum==1.16.1) (2023.12.2)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.10/site-packages (from huggingface-hub>=0.8.0->optimum==1.16.1) (2.28.1)\n",
      "Requirement already satisfied: tqdm>=4.42.1 in /opt/conda/lib/python3.10/site-packages (from huggingface-hub>=0.8.0->optimum==1.16.1) (4.66.1)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /opt/conda/lib/python3.10/site-packages (from huggingface-hub>=0.8.0->optimum==1.16.1) (6.0)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /opt/conda/lib/python3.10/site-packages (from huggingface-hub>=0.8.0->optimum==1.16.1) (4.9.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /opt/conda/lib/python3.10/site-packages (from packaging->optimum==1.16.1) (3.0.9)\n",
      "Requirement already satisfied: networkx in /opt/conda/lib/python3.10/site-packages (from torch>=1.9->optimum==1.16.1) (2.8.5)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.10/site-packages (from torch>=1.9->optimum==1.16.1) (3.1.2)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch>=1.9->optimum==1.16.1) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch>=1.9->optimum==1.16.1) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch>=1.9->optimum==1.16.1) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /opt/conda/lib/python3.10/site-packages (from torch>=1.9->optimum==1.16.1) (8.9.2.26)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /opt/conda/lib/python3.10/site-packages (from torch>=1.9->optimum==1.16.1) (12.1.3.1)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /opt/conda/lib/python3.10/site-packages (from torch>=1.9->optimum==1.16.1) (11.0.2.54)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /opt/conda/lib/python3.10/site-packages (from torch>=1.9->optimum==1.16.1) (10.3.2.106)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /opt/conda/lib/python3.10/site-packages (from torch>=1.9->optimum==1.16.1) (11.4.5.107)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /opt/conda/lib/python3.10/site-packages (from torch>=1.9->optimum==1.16.1) (12.1.0.106)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.19.3 in /opt/conda/lib/python3.10/site-packages (from torch>=1.9->optimum==1.16.1) (2.19.3)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch>=1.9->optimum==1.16.1) (12.1.105)\n",
      "Requirement already satisfied: triton==2.2.0 in /opt/conda/lib/python3.10/site-packages (from torch>=1.9->optimum==1.16.1) (2.2.0)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12 in /opt/conda/lib/python3.10/site-packages (from nvidia-cusolver-cu12==11.4.5.107->torch>=1.9->optimum==1.16.1) (12.3.101)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /opt/conda/lib/python3.10/site-packages (from transformers>=4.26.0->transformers[sentencepiece]>=4.26.0->optimum==1.16.1) (2022.7.9)\n",
      "Requirement already satisfied: tokenizers<0.19,>=0.14 in /opt/conda/lib/python3.10/site-packages (from transformers>=4.26.0->transformers[sentencepiece]>=4.26.0->optimum==1.16.1) (0.15.2)\n",
      "Requirement already satisfied: safetensors>=0.3.1 in /opt/conda/lib/python3.10/site-packages (from transformers>=4.26.0->transformers[sentencepiece]>=4.26.0->optimum==1.16.1) (0.4.1)\n",
      "Requirement already satisfied: sentencepiece!=0.1.92,>=0.1.91 in /opt/conda/lib/python3.10/site-packages (from transformers[sentencepiece]>=4.26.0->optimum==1.16.1) (0.1.96)\n",
      "Requirement already satisfied: protobuf in /opt/conda/lib/python3.10/site-packages (from transformers[sentencepiece]>=4.26.0->optimum==1.16.1) (3.20.3)\n",
      "Requirement already satisfied: humanfriendly>=9.1 in /opt/conda/lib/python3.10/site-packages (from coloredlogs->optimum==1.16.1) (10.0)\n",
      "Requirement already satisfied: pyarrow>=6.0.0 in /opt/conda/lib/python3.10/site-packages (from datasets->optimum==1.16.1) (8.0.0)\n",
      "Requirement already satisfied: dill<0.3.6 in /opt/conda/lib/python3.10/site-packages (from datasets->optimum==1.16.1) (0.3.5.1)\n",
      "Requirement already satisfied: pandas in /opt/conda/lib/python3.10/site-packages (from datasets->optimum==1.16.1) (1.4.3)\n",
      "Requirement already satisfied: xxhash in /opt/conda/lib/python3.10/site-packages (from datasets->optimum==1.16.1) (3.0.0)\n",
      "Requirement already satisfied: multiprocess in /opt/conda/lib/python3.10/site-packages (from datasets->optimum==1.16.1) (0.70.13)\n",
      "Requirement already satisfied: aiohttp in /opt/conda/lib/python3.10/site-packages (from datasets->optimum==1.16.1) (3.8.1)\n",
      "Requirement already satisfied: responses<0.19 in /opt/conda/lib/python3.10/site-packages (from datasets->optimum==1.16.1) (0.18.0)\n",
      "Requirement already satisfied: mpmath>=0.19 in /opt/conda/lib/python3.10/site-packages (from sympy->optimum==1.16.1) (1.3.0)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /opt/conda/lib/python3.10/site-packages (from aiohttp->datasets->optimum==1.16.1) (23.2.0)\n",
      "Requirement already satisfied: charset-normalizer<3.0,>=2.0 in /opt/conda/lib/python3.10/site-packages (from aiohttp->datasets->optimum==1.16.1) (2.1.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /opt/conda/lib/python3.10/site-packages (from aiohttp->datasets->optimum==1.16.1) (6.0.2)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /opt/conda/lib/python3.10/site-packages (from aiohttp->datasets->optimum==1.16.1) (4.0.2)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /opt/conda/lib/python3.10/site-packages (from aiohttp->datasets->optimum==1.16.1) (1.7.2)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /opt/conda/lib/python3.10/site-packages (from aiohttp->datasets->optimum==1.16.1) (1.3.0)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /opt/conda/lib/python3.10/site-packages (from aiohttp->datasets->optimum==1.16.1) (1.2.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.10/site-packages (from requests->huggingface-hub>=0.8.0->optimum==1.16.1) (3.3)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/conda/lib/python3.10/site-packages (from requests->huggingface-hub>=0.8.0->optimum==1.16.1) (1.26.10)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.10/site-packages (from requests->huggingface-hub>=0.8.0->optimum==1.16.1) (2022.6.15)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.10/site-packages (from jinja2->torch>=1.9->optimum==1.16.1) (2.1.3)\n",
      "Requirement already satisfied: python-dateutil>=2.8.1 in /opt/conda/lib/python3.10/site-packages (from pandas->datasets->optimum==1.16.1) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in /opt/conda/lib/python3.10/site-packages (from pandas->datasets->optimum==1.16.1) (2022.1)\n",
      "Requirement already satisfied: six>=1.5 in /opt/conda/lib/python3.10/site-packages (from python-dateutil>=2.8.1->pandas->datasets->optimum==1.16.1) (1.16.0)\n",
      "Requirement already satisfied: transformers==4.36.0 in /opt/conda/lib/python3.10/site-packages (4.36.0)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.10/site-packages (from transformers==4.36.0) (3.13.1)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.19.3 in /opt/conda/lib/python3.10/site-packages (from transformers==4.36.0) (0.20.1)\n",
      "Requirement already satisfied: numpy>=1.17 in /opt/conda/lib/python3.10/site-packages (from transformers==4.36.0) (1.23.4)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.10/site-packages (from transformers==4.36.0) (21.3)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /opt/conda/lib/python3.10/site-packages (from transformers==4.36.0) (6.0)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /opt/conda/lib/python3.10/site-packages (from transformers==4.36.0) (2022.7.9)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.10/site-packages (from transformers==4.36.0) (2.28.1)\n",
      "Requirement already satisfied: tokenizers<0.19,>=0.14 in /opt/conda/lib/python3.10/site-packages (from transformers==4.36.0) (0.15.2)\n",
      "Requirement already satisfied: safetensors>=0.3.1 in /opt/conda/lib/python3.10/site-packages (from transformers==4.36.0) (0.4.1)\n",
      "Requirement already satisfied: tqdm>=4.27 in /opt/conda/lib/python3.10/site-packages (from transformers==4.36.0) (4.66.1)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /opt/conda/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.19.3->transformers==4.36.0) (2023.12.2)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /opt/conda/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.19.3->transformers==4.36.0) (4.9.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /opt/conda/lib/python3.10/site-packages (from packaging>=20.0->transformers==4.36.0) (3.0.9)\n",
      "Requirement already satisfied: charset-normalizer<3,>=2 in /opt/conda/lib/python3.10/site-packages (from requests->transformers==4.36.0) (2.1.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.10/site-packages (from requests->transformers==4.36.0) (3.3)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/conda/lib/python3.10/site-packages (from requests->transformers==4.36.0) (1.26.10)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.10/site-packages (from requests->transformers==4.36.0) (2022.6.15)\n"
     ]
    }
   ],
   "source": [
    "!pip install optimum==1.16.1\n",
    "!pip install transformers==4.36.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b0471b34-bdaa-42f1-af0f-1468c560b0fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "  _torch_pytree._register_pytree_node(\n",
      "2024-02-16 16:01:44.185209: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2024-02-16 16:01:44.185256: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2024-02-16 16:01:44.186646: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "/opt/conda/lib/python3.10/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "  _torch_pytree._register_pytree_node(\n",
      "/opt/conda/lib/python3.10/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "  _torch_pytree._register_pytree_node(\n",
      "/opt/conda/lib/python3.10/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "  _torch_pytree._register_pytree_node(\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModelForSequenceClassification, AutoTokenizer\n",
    "from optimum.onnxruntime import ORTModelForSequenceClassification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9c22987b-f4a6-415e-9a40-327d9b69a2fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# a pre-trained model from HuggingFace\n",
    "model_name = \"bert-base-uncased\"\n",
    "onnx_directory = \"bert-base-uncased_onnx\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6428193d-f0e4-4ad8-b140-6f5fe797dd65",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "# Load the pre-trained model and tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\", export=True)\n",
    "model = AutoModelForSequenceClassification.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fb324d63-3bee-41ed-95bd-40bf49b13d63",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SequenceClassifierOutput(loss=None, logits=tensor([[-0.0306,  0.2759]], grad_fn=<AddmmBackward0>), hidden_states=None, attentions=None)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_ids = tokenizer(\"I love PyTorch!\", return_tensors=\"pt\")\n",
    "model(**input_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "62746dc3-f43c-44eb-8c4c-c17283d26a94",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Framework not specified. Using pt to export to ONNX.\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "Using the export variant default. Available variants are:\n",
      "    - default: The default ONNX variant.\n",
      "Using framework PyTorch: 2.2.0+cu121\n",
      "Overriding 1 configuration item(s)\n",
      "\t- use_cache -> False\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "SequenceClassifierOutput(loss=None, logits=tensor([[ 0.1623, -0.0475]]), hidden_states=None, attentions=None)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_onnx = ORTModelForSequenceClassification.from_pretrained(\"bert-base-uncased\", export=True)\n",
    "model_onnx(**input_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1f88c2bc-4bea-4fa0-8487-ad406e94719d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('bert-base-uncased_onnx/tokenizer_config.json',\n",
       " 'bert-base-uncased_onnx/special_tokens_map.json',\n",
       " 'bert-base-uncased_onnx/vocab.txt',\n",
       " 'bert-base-uncased_onnx/added_tokens.json',\n",
       " 'bert-base-uncased_onnx/tokenizer.json')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_onnx.save_pretrained(onnx_directory)\n",
    "tokenizer.save_pretrained(onnx_directory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4df44901-4dd0-4670-a348-94a027e63419",
   "metadata": {},
   "outputs": [],
   "source": [
    "from optimum.onnxruntime.configuration import AutoQuantizationConfig\n",
    "from optimum.onnxruntime import ORTQuantizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "668a7a50-daf8-4648-bc54-720643f0411a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Creating dynamic quantizer: QOperator (mode: IntegerOps, schema: u8/s8, channel-wise: False)\n",
      "Quantizing model...\n",
      "Saving quantized model at: bert-base-uncased_onnx (external data format: False)\n",
      "Configuration saved in bert-base-uncased_onnx/ort_config.json\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "PosixPath('bert-base-uncased_onnx')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qconfig = AutoQuantizationConfig.arm64(is_static=False, per_channel=False)\n",
    "quantizer = ORTQuantizer.from_pretrained(model_onnx)\n",
    "quantizer.quantize(save_dir=onnx_directory, quantization_config=qconfig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "92b1f6a7-bf31-48d3-aabf-bdea6aa617d0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SequenceClassifierOutput(loss=None, logits=tensor([[ 0.1623, -0.0475]]), hidden_states=None, attentions=None)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_quantized = ORTModelForSequenceClassification.from_pretrained(onnx_directory, file_name=\"model.onnx\")\n",
    "model_quantized(**input_ids)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (Local)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
